# # 提取DCT块不同频率系数，并且合并
# image = torch.rand(c, 3, w, h)
# image1 = image.reshape(c, int(3*w*h/64), 8, 8)
# image1 = image1.reshape(c, 3, int(w*h/64), 64)
# image1 = image1.permute(0, 1, 3, 2)
# image1 = image1.reshape(c, 3, w, h)
# # 逆处理，还原为原来的DCT
# ima = image1.reshape(c, 3, 64, int(w*h/64))
# ima = ima.permute(0, 1, 3, 2)
# ima = ima.reshape(c, int(3*w*h/64), 8, 8)
# ima = ima.reshape(c, 3, w, h)

import torch
import torch.nn as nn
import torch.nn.functional as F

# x = torch.rand(2, 3, 16, 16)
# print(x.shape)
# x = x.reshape(2, 3, int(16 * 16 / 64), 64)
# x = x.permute(0, 1, 3, 2)
# x = x.reshape(2, int(3 * 8 * 8), 2, 2)
#
# y = x.reshape(2, 3, 64, 4)
# y = y.permute(0, 1, 3, 2)
# y = y.reshape(2, 3, 16, 16)
# print(y.shape)
# print(y.equal(x))
# class MaxpoolNet(nn.Module):
#     def __init__(self):
#         super(MaxpoolNet, self).__init__()
#         self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
#         self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
#         self.Maxsample = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
#         self.conv3 = nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1)
#
#     def forward(self, x):
#         x = F.relu(self.conv1(x))
#         x = F.relu(self.conv2(x))
#         x = self.Maxsample(x)
#         x = torch.sigmoid(self.conv3(x))
#         return x
#
# # Create an instance of the network
# model = MaxpoolNet()
# # Create a random input tensor of size 1x1x120x120
# x = torch.randn(32, 1, 256, 256)
# # Upsample the input to size 1x1x256x256
# y = model(x)
# print(y.shape)

image = torch.rand(1,3,16,16)
print(image)
message = torch.rand(16, 16)
message = message.reshape(1, 1, 16, 16)
image = torch.cat([image, message], dim=1)
print(image)



"""
    DCT频率提取与还原
"""
# 频率提取
def DCT2frequency(image, c, w):
    image = image.reshape(c, 3, int(w*w/64), 64)
    image = image.permute(0, 1, 3, 2)
    image = image.reshape(c, 3*64, int(w/8), int(w/8))
    return image

# 逆操作
def frequency2DCT(image, c, w):
    ima = image.reshape(c, 3, 64, int(w * w / 64))
    ima = ima.permute(0, 1, 3, 2)
    ima = ima.reshape(c, 3, w, w)
    return ima